module stationary_vortex_test_mod

! reference: Nair and Jablonowski, 2008. Moving Vortices on the Sphere: A Test Case for Horizontal Advection Problems. Mon.Wea.Rev
	use const_mod
	use sphere_geometry_mod
	use mesh_mod
	use state_mod
	use namelist_mod
	use tracer_mod
	use parallel_mod

	implicit none

	private
	public stationary_vortex_test_init
	public stationary_vortex_test_set_ic
	public stationary_vortex_test_get_exact
	public stationary_vortex_test_set_uv

	real(r8), parameter :: period = 12 * 86400
	real(r8), parameter :: v0     = pi2 * radius / period
  real(r8) :: lonp, latp, rho0, gamma
contains

	subroutine stationary_vortex_test_init()

		rho0 = 3.0
		lonp = 0; latp = pi05
		gamma = 5

	end subroutine stationary_vortex_test_init

	subroutine stationary_vortex_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, itracer
    real(r8) lon, lat, lon_rot, lat_rot
    
    if (ntracers /= 1) then
    	print*, 'stationary_vortex_test case only need 1 tracers!'
    	stop
    end if

    do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    	lat = mesh%full_lat(j)
    	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    		lon = mesh%full_lon(i)
    		call regular2rotate(lon, lat, lonp, latp, lon_rot, lat_rot)
    		tracer%q(i,j,1,1) = 1 - tanh(rho0 * cos(lat_rot) / gamma * sin(lon_rot))
    	end do
    end do

    do itracer = 1, size(tracer%q, 4)
    	call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    	call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    end do

	end subroutine stationary_vortex_test_set_ic

	subroutine stationary_vortex_test_get_exact(mesh, tracer, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		real(8), intent(in) :: time_in_seconds
		integer i, j, itracer
    real(r8) lon, lat, lon_rot, lat_rot
    
    if (ntracers > 1) then
    	print*, 'stationary_vortex_test case only need 1 tracers!'
    	stop
    end if

    do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    	lat = mesh%full_lat(j)
    	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    		lon = mesh%full_lon(i)
    		call regular2rotate(lon, lat, lonp, latp, lon_rot, lat_rot)
    		tracer%qanl(i,j) = 1 - tanh(rho0 * cos(lat_rot) / gamma * sin(lon_rot - romega(lat_rot) / radius * time_in_seconds))
    	end do
    end do

	end subroutine stationary_vortex_test_get_exact

	subroutine stationary_vortex_test_set_uv(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds

		real(r8) lon, lat, lon_rot, lat_rot
		integer i, j

		associate (u => state%u, v => state%v)

			do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
				lat = mesh%full_lat(j)
				do i = mesh%half_lon_ibeg, mesh%half_lon_iend
					lon = mesh%half_lon(i)
					call regular2rotate(lon, lat, lonp, latp, lon_rot, lat_rot)
					u(i,j,1) = romega(lat_rot) * (sin(latp) * cos(lat) - cos(latp) * cos(lon - lonp) * sin(lat))
				end do
			end do
      call fill_zonal_halo_lon(u, all_halo=.true.)

			do j = mesh%half_lat_ibeg, mesh%half_lat_iend
				lat = mesh%half_lat(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					lon = mesh%full_lon(i)
					call regular2rotate(lon, lat, lonp, latp, lon_rot, lat_rot)
					v(i,j,1) = romega(lat_rot) * cos(latp) * sin(lon - lonp)
				end do
			end do

		end associate

	end subroutine stationary_vortex_test_set_uv

	pure real(r8) function romega(lat) result(res)

		real(r8), intent(in) :: lat
		real(r8) rho

		rho = rho0 * cos(lat)
		if (abs(rho) < eps) then 
			res = 0.
		else
			res = v0 * 1.5 * sqrt(3.0) * tanh(rho) / cosh(rho)**2 / rho
		end if

	end function romega

end module stationary_vortex_test_mod